from __future__ import annotations

from model import (
    Settings,
    Request,
    Action,
    ReloadAction,
    VisitAction,
    OffloadAction,
    FinAction,
)

import matplotlib
import matplotlib.pyplot as plt

matplotlib.use("Agg")


class Visualizer:
    """
    A class to visualize settings, requests, and actions.

    The settings, requests and actions must be validated before being passed to this class.
    """

    def __init__(
        self, settings: Settings, requests: list[Request], actions: list[Action]
    ) -> None:
        self.settings = settings
        self.requests = requests
        self.actions = actions

    @staticmethod
    def _assign_lanes(intervals: list[tuple]) -> dict:
        """
        Assign lanes (y positions) to intervals to avoid overlaps.
        intervals: list of tuples (start, end, ...)
        Returns: dict mapping interval index to lane number
        """
        if not intervals:
            return {}

        # Sort by start time, then by end time
        indexed_intervals = [
            (i, intervals[i][0], intervals[i][1]) for i in range(len(intervals))
        ]
        indexed_intervals.sort(key=lambda x: (x[1], x[2]))

        # Greedy lane assignment
        lanes = []  # Each lane stores the end time of the last interval in that lane
        lane_assignment = {}

        for idx, start, end in indexed_intervals:
            # Find the first available lane
            assigned = False
            for lane_num, lane_end in enumerate(lanes):
                if start >= lane_end:  # No overlap
                    lanes[lane_num] = end
                    lane_assignment[idx] = lane_num
                    assigned = True
                    break

            if not assigned:
                # Create a new lane
                lanes.append(end)
                lane_assignment[idx] = len(lanes) - 1

        return lane_assignment

    def visualize(self, path: str) -> None:
        """
        Visualize the given (validated) settings, requests, and actions.
        """
        # Build intervals for actions
        visits = []  # tuples (start, end, index)
        ios = []  # tuples (start, end, type, addr, size)
        for action in self.actions:
            if isinstance(action, VisitAction):
                start = action.time
                end = action.time + action.cost()
                visits.append((start, end, action.index))
            elif isinstance(action, (ReloadAction, OffloadAction)):
                start = action.time
                end = action.time + action.cost()
                typ = "Reload" if isinstance(action, ReloadAction) else "Offload"
                ios.append((start, end, typ, action.addr, action.size))
            elif isinstance(action, FinAction):
                # FinAction marks the reported finish time; include as vertical marker later
                fin_time = action.time

        # Sort intervals
        visits.sort()
        ios.sort()

        # Assign lanes for parallel display
        # Prepare intervals for requests
        request_intervals = [
            (req.start, req.start + req.time, i) for i, req in enumerate(self.requests)
        ]
        request_lanes = self._assign_lanes(request_intervals)
        num_request_lanes = (
            max(request_lanes.values(), default=0) + 1 if request_lanes else 1
        )

        # Assign lanes for visits
        visit_lanes = self._assign_lanes(visits)
        num_visit_lanes = max(visit_lanes.values(), default=0) + 1 if visit_lanes else 1

        # Assign lanes for IOs
        io_lanes = self._assign_lanes(ios)
        num_io_lanes = max(io_lanes.values(), default=0) + 1 if io_lanes else 1

        # Calculate figure height: each lane gets 0.4 inches, plus 3 inches for memory chart
        total_height = max(
            4, (num_request_lanes + num_visit_lanes + num_io_lanes) * 0.4 + 3
        )
        fig, axes = plt.subplots(
            4,
            1,
            figsize=(14, total_height),
            sharex=True,
            gridspec_kw={
                "height_ratios": [num_request_lanes, num_visit_lanes, num_io_lanes, 3]
            },
        )

        # Bar 1 (top): Requests - show the time window [start, start+time] for each request
        ax0 = axes[0]
        for i, req in enumerate(self.requests):
            start = req.start
            end = req.start + req.time
            lane = request_lanes.get(i, 0)
            # Use different colors cycling through a palette
            color = f"C{i % 10}"
            ax0.barh(
                lane,
                end - start,
                left=start,
                height=0.8,
                color=color,
                alpha=0.7,
                edgecolor="black",
                linewidth=1.5,
            )

        ax0.set_yticks(range(num_request_lanes))
        ax0.set_yticklabels([f"Lane {i}" for i in range(num_request_lanes)])
        ax0.set_ylabel("Requests")
        ax0.set_ylim(-0.5, num_request_lanes - 0.5)
        ax0.grid(True, axis="x", linestyle="--", alpha=0.3)

        # Bar 2 (middle): Visits - show execution time for each VisitAction
        ax1 = axes[1]
        for i, (s, e, idx) in enumerate(visits):
            lane = visit_lanes.get(i, 0)
            color = f"C{idx % 10}"
            ax1.barh(
                lane,
                e - s,
                left=s,
                height=0.8,
                color=color,
                alpha=0.6,
                edgecolor="black",
                linewidth=2,
            )

        ax1.set_yticks(range(num_visit_lanes))
        ax1.set_yticklabels([f"Lane {i}" for i in range(num_visit_lanes)])
        ax1.set_ylabel("Visits")
        ax1.set_ylim(-0.5, num_visit_lanes - 0.5)
        ax1.grid(True, axis="x", linestyle="--", alpha=0.3)

        # Bar 3 (bottom): IO - show execution time for each Reload/Offload
        ax2 = axes[2]
        for i, (s, e, typ, addr, size) in enumerate(ios):
            lane = io_lanes.get(i, 0)
            # Use distinct colors and hatch patterns for Reload vs Offload
            if typ == "Reload":
                color = "tab:blue"
                hatch = "///"  # Diagonal lines for Reload
            else:  # Offload
                color = "tab:red"
                hatch = "\\\\\\"  # Reverse diagonal lines for Offload

            ax2.barh(
                lane,
                e - s,
                left=s,
                height=0.8,
                color=color,
                alpha=0.5,
                edgecolor="black",
                linewidth=1.5,
                hatch=hatch,
            )

        ax2.set_yticks(range(num_io_lanes))
        ax2.set_yticklabels([f"Lane {i}" for i in range(num_io_lanes)])
        ax2.set_ylabel("IO")
        ax2.set_ylim(-0.5, num_io_lanes - 0.5)
        ax2.grid(True, axis="x", linestyle="--", alpha=0.3)

        # Bar 4 (bottom): Memory Usage - line chart showing memory occupation over time
        ax3 = axes[3]

        # Build memory usage timeline
        memory_events = []  # (time, delta_memory, event_type)

        for action in self.actions:
            if isinstance(action, ReloadAction):
                # Reload starts: memory increases instantly
                memory_events.append((action.time, action.size, "reload_start"))
            elif isinstance(action, OffloadAction):
                # Offload ends: memory decreases instantly
                offload_end_time = action.time + action.cost()
                memory_events.append((offload_end_time, -action.size, "offload_end"))

        # Sort events by time
        memory_events.sort(key=lambda x: x[0])

        # Build timeline with step function
        times = [0]
        memory_usage = [0]
        current_memory = 0

        for time, delta, event_type in memory_events:
            # Add point just before the change
            if times[-1] != time:
                times.append(time)
                memory_usage.append(current_memory)

            # Apply the change
            current_memory += delta
            times.append(time)
            memory_usage.append(current_memory)

        # Extend to fin_time
        if times[-1] < fin_time:
            times.append(fin_time)
            memory_usage.append(current_memory)

        # Plot memory usage
        ax3.plot(
            times, memory_usage, linewidth=2, color="tab:green", label="Memory Usage"
        )
        ax3.fill_between(times, 0, memory_usage, alpha=0.3, color="tab:green")

        # Add memory limit line
        ax3.axhline(
            self.settings.M,
            color="tab:red",
            linestyle="--",
            linewidth=1.5,
            label=f"Memory Limit ({self.settings.M})",
        )

        ax3.set_ylabel("Memory Usage (bytes)")
        ax3.set_xlabel("Time")
        ax3.grid(True, linestyle="--", alpha=0.3)
        ax3.legend(loc="upper left", bbox_to_anchor=(1.02, 1), fontsize=8, frameon=True)
        ax3.set_ylim(bottom=0)

        # Mark Fin time on all axes
        for ax in axes:
            ax.axvline(fin_time, color="k", linestyle=":", linewidth=2, label="Fin")

        # Add legend for IO types
        from matplotlib.patches import Patch

        legend_elements = [
            Patch(
                facecolor="tab:blue",
                edgecolor="black",
                hatch="///",
                alpha=0.5,
                label="Reload",
            ),
            Patch(
                facecolor="tab:red",
                edgecolor="black",
                hatch="\\\\\\",
                alpha=0.5,
                label="Offload",
            ),
        ]
        ax2.legend(
            handles=legend_elements,
            loc="upper left",
            bbox_to_anchor=(1.02, 1),
            fontsize=8,
            frameon=True,
        )

        fig.tight_layout(rect=[0, 0, 0.85, 1])  # Leave space on the right for legends
        fig.savefig(path)
        plt.close(fig)

        return None
